package com.shujia.flink.state

import org.apache.flink.api.common.eventtime.WatermarkStrategy
import org.apache.flink.api.common.functions.RuntimeContext
import org.apache.flink.api.common.serialization.SimpleStringSchema
import org.apache.flink.api.common.state.{ListState, ListStateDescriptor}
import org.apache.flink.configuration.Configuration
import org.apache.flink.connector.kafka.source.KafkaSource
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer
import org.apache.flink.streaming.api.functions.KeyedProcessFunction
import org.apache.flink.streaming.api.scala._
import org.apache.flink.util.Collector

import java.{lang, util}

object Demo4ListState {
  def main(args: Array[String]): Unit = {

    /**
     * 实时统计每个班级的平均年龄
     */

    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment

    val source: KafkaSource[String] = KafkaSource
      .builder[String]
      //kafka 集群列表
      .setBootstrapServers("master:9092,node1:9092,node2:9092")
      //消费的topic
      .setTopics("student")
      //消费者组
      .setGroupId("my-group")
      //读取数据的位置，earliest：从最早读取数据，latest：读取最新数据
      .setStartingOffsets(OffsetsInitializer.latest)
      .setValueOnlyDeserializer(new SimpleStringSchema())
      .build

    //使用kafka 数据源
    val studentDS: DataStream[String] = env
      .fromSource(source, WatermarkStrategy.noWatermarks(), "Kafka Source")

    //取出班级和年龄
    val clazzAndAgeDS: DataStream[(String, Int)] = studentDS.map(stu => {
      val split: Array[String] = stu.split(",")
      val clazz: String = split(4)
      val age: Int = split(2).toInt
      (clazz, age)
    })

    //按照班级分组
    val keyByDS: KeyedStream[(String, Int), String] = clazzAndAgeDS.keyBy(_._1)

    //计算平均年龄
    val clazzAvgAgeDS: DataStream[(String, Double)] = keyByDS
      .process(new KeyedProcessFunction[String, (String, Int), (String, Double)] {

        //集合状态，相当于java中的集合
        var agesState: ListState[Int] = _

        override def open(parameters: Configuration): Unit = {
          //获取flink运行环境
          val context: RuntimeContext = getRuntimeContext

          //创建状态描述对象
          val stateDesc = new ListStateDescriptor[Int]("ages", classOf[Int])

          //创建状态
          agesState = context.getListState(stateDesc)
        }

        override def processElement(kv: (String, Int),
                                    ctx: KeyedProcessFunction[String, (String, Int), (String, Double)]#Context,
                                    out: Collector[(String, Double)]): Unit = {

          //清除状态中保存的数据
          //agesState.clear()

          val (clazz: String, age: Int) = kv

          //将年龄保存到集合状态中
          agesState.add(age)

          //计算平均年龄

          //获取状态周昂保存的年龄,返回一个迭代器
          val ages: util.Iterator[Int] = agesState.get().iterator()

          var sumAge = 0.0
          var num = 0

          //循环计算总人数和总年龄
          while (ages.hasNext) {
            sumAge += ages.next()
            num += 1
          }

          //平均年龄
          val avgAge: Double = sumAge / num

          //将结果发生到下游
          out.collect((clazz, avgAge))
        }
      })

    clazzAvgAgeDS.print()

    env.execute()

  }

}
